import pandas as pd
import pdfplumber
from docx import Document
import io


def process_pdf(filepath):
    """处理包含表格的PDF文档"""
    full_text = []
    with pdfplumber.open(filepath) as pdf:
        for page in pdf.pages:
            # 提取文本
            text = page.extract_text()
            if text:
                full_text.append(text)

            # 提取表格
            for table in page.extract_tables():
                df = pd.DataFrame(table[1:], columns=table[0])
                full_text.append("\n表格内容：")
                full_text.append(df.to_markdown(index=False))

    return "\n".join(full_text)


def process_docx(filepath):
    """处理包含公式的Word文档"""
    doc = Document(filepath)
    full_text = []

    for para in doc.paragraphs:
        full_text.append(para.text)

    # 处理表格
    for table in doc.tables:
        table_text = []
        for row in table.rows:
            row_text = [cell.text for cell in row.cells]
            table_text.append(" | ".join(row_text))
        full_text.append("\n表格：")
        full_text.append("\n".join(table_text))

    return "\n".join(full_text)


def process_excel(filepath):
    """处理Excel文件"""
    full_text = []
    xls = pd.ExcelFile(filepath)

    for sheet_name in xls.sheet_names:
        df = pd.read_excel(xls, sheet_name=sheet_name)
        full_text.append(f"\n工作表 [{sheet_name}]：")
        full_text.append(df.to_markdown(index=False))

    return "\n".join(full_text)


def read_file_content(filepath, extension):
    """根据文件类型调用不同处理器"""
    try:
        if extension == 'pdf':
            return process_pdf(filepath)
        elif extension == 'docx':
            return process_docx(filepath)
        elif extension in ['xlsx', 'xls']:
            return process_excel(filepath)
        elif extension == 'txt':
            with open(filepath, 'r', encoding='utf-8') as f:
                return f.read()
        else:
            raise ValueError("不支持的文件类型")
    except Exception as e:
        raise RuntimeError(f"文件解析失败: {str(e)}")


# 修改process_uploaded_file函数
def process_uploaded_file(file):
    try:
        filename = validate_file(file)
        filepath = os.path.join(
            current_app.config['UPLOAD_FOLDER'],
            filename
        )
        file.save(filepath)

        # 获取文件扩展名
        ext = filename.rsplit('.', 1)[1].lower()

        # 调用文件处理器
        content = read_file_content(filepath, ext)

        os.remove(filepath)
        return content
    except Exception as e:
        current_app.logger.error(f"文件处理失败: {str(e)}")
        raise